import cv2
import numpy as np
from paddleocr import PaddleOCR
from typing import Tuple, Optional

class PlateRecognizer:
    def __init__(self):
        self.ocr = PaddleOCR(use_angle_cls=True, lang="ch", use_gpu=False)

    def recognize_plate(self, image_input) -> Tuple[Optional[str], Optional[float]]:
        """
        车牌识别主函数，支持多种输入类型。

        参数:
            image_input: 可以是图像路径（str）、NumPy 数组或 OpenCV 图像对象。

        返回:
            tuple: (识别结果文本, 置信度) 或 (None, None) 如果识别失败
        """
        try:
            # 1. 处理输入
            if isinstance(image_input, str):  # 输入为图像路径
                img = cv2.imread(image_input)
            elif isinstance(image_input, np.ndarray):  # 输入为 NumPy 数组
                img = image_input
            elif hasattr(image_input, 'shape'):  # 输入为 OpenCV 图像对象
                img = image_input
            else:
                print("不支持的输入类型，请输入图像路径、NumPy 数组或 OpenCV 图像对象。")
                return None, None

            # 2. 检查图像是否有效
            if img is None or img.size == 0:
                print("输入的图像无效或为空。")
                return None, None

            # 3. 预处理和识别
            corrected_img = self._detect_and_correct_plate(img)
            if corrected_img is None:
                print("车牌检测失败，尝试直接识别")
                result = self.ocr.ocr(img, cls=True, det=False)
            else:
                # 对校正后的图像进行预处理
                gray = cv2.cvtColor(corrected_img, cv2.COLOR_BGR2GRAY)
                equalized = cv2.equalizeHist(gray)
                denoised = cv2.fastNlMeansDenoising(equalized, None, h=10,
                                                    templateWindowSize=7, searchWindowSize=21)
                result = self.ocr.ocr(denoised, cls=True, det=False)

            # 4. 解析结果
            if result and result[0]:
                text = result[0][0][0]
                confidence = result[0][0][1]
                return text, confidence
            return None, None

        except Exception as e:
            print(f"识别过程中出错: {str(e)}")
            return None, None

    def _detect_and_correct_plate(self, img):
        """
        检测并校正车牌图像。

        参数:
            img: 输入的图像数组。

        返回:
            校正后的图像数组，如果检测失败则返回 None。
        """
        try:
            # 使用 PaddleOCR 检测车牌区域（只检测不识别）
            result = self.ocr.ocr(img, cls=True, rec=False)

            if not result or not result[0]:
                return None

            # 获取检测到的最大区域（假设是车牌）
            max_area = 0
            plate_box = None

            for box in result[0]:
                try:
                    if len(box) != 4:
                        continue

                    box_np = np.array(box, dtype=np.int32)
                    area = cv2.contourArea(box_np)
                    if area > max_area:
                        max_area = area
                        plate_box = box
                except:
                    continue

            if plate_box is None:
                return None

            # 透视校正
            return self._correct_perspective(img, plate_box)

        except Exception as e:
            print(f"车牌检测和校正过程中出错: {str(e)}")
            return None

    def _correct_perspective(self, img, box):
        """
        透视校正车牌图像。

        参数:
            img: 输入的图像数组。
            box: 检测到的车牌区域框。

        返回:
            校正后的图像数组，如果校正失败则返回 None。
        """
        try:
            src_points = np.array(box, dtype=np.float32)
            if len(src_points) != 4:
                return None

            # 中国车牌标准比例
            target_ratio = 3.14
            width = np.linalg.norm(src_points[1] - src_points[0])
            height = np.linalg.norm(src_points[3] - src_points[0])
            current_ratio = width / height

            if current_ratio > target_ratio:
                target_height = int(width / target_ratio)
                target_width = int(width)
            else:
                target_width = int(height * target_ratio)
                target_height = int(height)

            dst_points = np.array([
                [0, 0],
                [target_width - 1, 0],
                [target_width - 1, target_height - 1],
                [0, target_height - 1]
            ], dtype=np.float32)

            # 对源点进行排序
            rect = np.zeros((4, 2), dtype=np.float32)
            s = src_points.sum(axis=1)
            rect[0] = src_points[np.argmin(s)]
            rect[2] = src_points[np.argmax(s)]
            diff = np.diff(src_points, axis=1)
            rect[1] = src_points[np.argmin(diff)]
            rect[3] = src_points[np.argmax(diff)]

            M = cv2.getPerspectiveTransform(rect, dst_points)
            return cv2.warpPerspective(img, M, (target_width, target_height))

        except Exception as e:
            print(f"透视校正过程中出错: {str(e)}")
            return None
